from keras import Input
from keras.layers import Conv2D

from .base_layer import BaseLayer
from config import ModelLayers
from entity.model_describe import ModelDescribe
from log import log


class Conv2DLayer(BaseLayer):
    def persistence(self, model_list_line: list, train_id: int) -> ModelDescribe:
        m = super().persistence(model_list_line, train_id)
        filters = model_list_line[1]
        kernel_size = model_list_line[2]
        m.var1 = filters
        m.var2 = kernel_size
        m.var3 = True if (len(model_list_line) == 4 and model_list_line[3] == True) else False
        return m

    def check(self, model_list_line: list, models: list = None) -> bool:
        if len(model_list_line) < 3:
            log.error("传入参数不足")
            raise RuntimeError("传入参数不足")
        return True

    def transfer(self, model_list_line: list, inputs: Input, models: list = None):

        filters = model_list_line[1]
        kernel_size = model_list_line[2]
        # 这是业务逻辑check不检查
        if len(model_list_line) == 4 and model_list_line[3] == True:
            inputs = Conv2D(filters, (kernel_size, kernel_size), padding='same', activation='relu')(inputs)
        else:
            inputs = Conv2D(filters, (kernel_size, kernel_size), padding='same')(inputs)
        return inputs
